package com.jboltai.capability.text.extract.pdf;

import cn.hutool.core.util.RandomUtil;
import cn.hutool.core.util.StrUtil;
import com.jboltai.capability.text.ContentInfo;
import com.jboltai.capability.text.ContentType;
import com.jboltai.event.text.TextEvent;
import com.jboltai.util.text.FileInfo;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.pdmodel.PDPageTree;
import org.apache.pdfbox.rendering.ImageType;
import org.apache.pdfbox.rendering.PDFRenderer;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.InputStream;
import java.util.List;
import java.util.Objects;
import java.util.function.Consumer;

public class PdfExtractor {


    public static void extract(File file, Consumer<String> chunkConsumer, TextEvent textEvent) {

        try (PDDocument document = PDDocument.load(file)) {

            if (document.isEncrypted()) {
                throw new RuntimeException("文档被加密，无法提取内容");
            }
            PDFRenderer pdfRenderer = new PDFRenderer(document);
            PDPageTree pages = document.getPages();
            TextStripper stripper = new TextStripper();
            ImageLocations imgProcessor = new ImageLocations();
            StringBuilder content = new StringBuilder();
            for (int i = 0; i < pages.getCount(); i++) {

                stripper.setStartPage(i + 1);
                stripper.setEndPage(i + 1);

                List<ContentInfo> contentInfos = stripper.readTextInfo(document);
                if (textEvent.getFile2TextConverter() != null) {
                    //设置了文件转换器才需要去提取图片
                    contentInfos = imgProcessor.processPage(pages.get(i), contentInfos);
                }
                //判断是否需要处理文件，如果设置了拦截器，要根据拦截器判断是否触发onfile
                PdfFileInterceptor interceptor = (PdfFileInterceptor) textEvent.getAttr(PdfFileInterceptor.INNER_KEY);
                boolean handleFile = textEvent.getFile2TextConverter() != null && (interceptor == null || interceptor.intercept((int)contentInfos.stream().filter(item -> item.getType() == ContentType.TEXT).count(),(int)contentInfos.stream().filter(item -> item.getType() != ContentType.TEXT).count()));
                for (ContentInfo contentInfo : contentInfos) {
                    switch (contentInfo.getType()) {
                        case TEXT:
                            if (StrUtil.isNotBlank(contentInfo.getText())) {
                                content.append(contentInfo.getText()).append("\n");
                            }
                            break;
                        case IMG: {

                            if (handleFile) {

                                FileInfo img = new FileInfo();
                                img.setFileName(RandomUtil.randomString(6) + "." + contentInfo.getImage().getSuffix());
                                try(InputStream is = convertBufferedImageToInputStream(contentInfo.getImage().getImage(), contentInfo.getImage().getSuffix())) {

                                    img.setInputStream(is);
                                    String convertText = textEvent.getFile2TextConverter().convert(textEvent, img);
//                                img.setConvertText(convertText);
                                    content.append(convertText).append("\n");
                                }


                            }
                            break;
                        }
                        case VIDEO:
                        case OTHER:
                            //todo 待完善
                            break;
                    }
                }

                if (textEvent.getFile2TextConverter() != null && interceptor != null && !handleFile && Objects.equals(textEvent.getAttr(PdfFileInterceptor.PAGE_TO_IMG_KEY), true)) {
                    //某一页的文件太多了，不需要处理，将这页转成图片
                    BufferedImage bufferedImage = pdfRenderer.renderImageWithDPI(i, 300, ImageType.RGB);
                    FileInfo img = new FileInfo();
                    img.setFileName(RandomUtil.randomString(6) + ".png" );
                    try(InputStream is = convertBufferedImageToInputStream(bufferedImage, "png")) {

                        img.setInputStream(is);
                        String convertText = textEvent.getFile2TextConverter().convert(textEvent, img);
//                                img.setConvertText(convertText);
                        content.append(convertText).append("\n");
                    }
                }


                //流式读取
                chunkConsumer.accept(content.toString());
                content.setLength(0);


            }


        } catch (Exception e) {
//            textEvent.fail(EventErrorType.UNKNOWN, e);
            throw new RuntimeException(e);
        } finally {
//            textEvent.complete();
        }

    }

    private static InputStream convertBufferedImageToInputStream(BufferedImage image, String formatName) {
        try {
            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
            ImageIO.write(image, formatName, outputStream);
            return new ByteArrayInputStream(outputStream.toByteArray());
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

}
